package com.gitee.zycra.file.utils;

import org.bouncycastle.crypto.BlockCipher;
import org.bouncycastle.crypto.engines.AESEngine;
import org.bouncycastle.crypto.engines.SM4Engine;
import org.bouncycastle.crypto.io.CipherInputStream;
import org.bouncycastle.crypto.modes.AEADBlockCipher;
import org.bouncycastle.crypto.modes.GCMBlockCipher;
import org.bouncycastle.crypto.params.AEADParameters;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;

/**
 * Symmetric encrypt util for AES and SM4.
 *
 * @author zycra
 * @since 1.0.0
 */
public class SymmetricEncryptUtil extends BCUtil {

    private static final Logger LOGGER = LoggerFactory.getLogger(SymmetricEncryptUtil.class);

    protected static final String SM4_LABEL = "SM4";

    protected static final String AES_LABEL = "AES";

    /**
     * Buffer size of read file.
     *
     * @since 1.0.0
     */
    private static final int BUFFER_SIZE = 4096;

    /**
     * Mac size of AEAD parameters.
     *
     * @since 1.0.0
     */
    private static final int MAC_SIZE = 128;

    /**
     * IV size of AEAD parameters.
     *
     * @since 1.0.0
     */
    private static final int IV_SIZE = 12;

    /**
     * SM4 key size.
     *
     * @since 1.0.0
     */
    private static final int SM4_KEY_SIZE = 128;

    /**
     * AES key size.
     *
     * @since 1.0.0
     */
    private static final int AES_KEY_SIZE = 256;

    /**
     * IV for create AEAD parameters.
     *
     * @since 1.0.0
     */
    private static final byte[] IV = new byte[IV_SIZE];

    /**
     * All hex characters.
     *
     * @since 1.0.0
     */
    private static final byte[] HEX_ARRAY = "0123456789ABCDEF".getBytes(StandardCharsets.UTF_8);

    private static final int TWO = 2;

    private static final int FOUR = 4;

    private static final int SIXTEEN = 16;

    private static final int F = 0x0F;

    /**
     * The key size map, key is encrypt key label and value is key size.
     *
     * @since 1.0.0
     */
    private static final Map<String, Integer> KEY_SIZE_MAP = new HashMap<>();

    /**
     * Block cipher size map, key is encrypt key label and value is {@link BlockCipher} instance.
     *
     * @since 1.0.0
     */
    private static final Map<String, BlockCipher> BLOCK_CIPHER_MAP = new HashMap<>();

    static {
        KEY_SIZE_MAP.put(SM4_LABEL, SM4_KEY_SIZE);
        BLOCK_CIPHER_MAP.put(SM4_LABEL, new SM4Engine());
        KEY_SIZE_MAP.put(AES_LABEL, AES_KEY_SIZE);
        BLOCK_CIPHER_MAP.put(AES_LABEL, new AESEngine());
    }

    /**
     * Encrypt content by password and key label.
     *
     * @param content  the content to encrypt.
     * @param password the password for encrypt.
     * @param label    the key label.
     * @return encrypt result.
     * @since 1.0.0
     */
    protected static String encrypt(String content, String password, String label) {
        return doEncrypt(content, password, true, label);
    }

    /**
     * Decrypt content by password and key label.
     *
     * @param content  the content to decrypt.
     * @param password the password for decrypt.
     * @param label    the key label.
     * @return decrypt result.
     * @since 1.0.0
     */
    protected static String decrypt(String content, String password, String label) {
        return doEncrypt(content, password, false, label);
    }

    /**
     * Encrypt file by password and key label.
     *
     * @param originFile  the origin file to encrypt.
     * @param encryptFile the target path to store encrypt file.
     * @param password    the password for encrypt.
     * @param label       the key label.
     * @return whether successful.
     * @since 1.0.0
     */
    protected static boolean encrypt(File originFile, File encryptFile, String password, String label) {
        return doEncryptFile(originFile, encryptFile, password, true, label);
    }

    /**
     * Decrypt file by password and key label.
     *
     * @param originFile  the origin file to decrypt.
     * @param decryptFile the target path to store decrypt file.
     * @param password    the password for decrypt.
     * @param label       the key label.
     * @return whether successful.
     * @since 1.0.0
     */
    protected static boolean decrypt(File originFile, File decryptFile, String password, String label) {
        return doEncryptFile(originFile, decryptFile, password, false, label);
    }

    /**
     * Generate a key by given key label.
     *
     * @param label the key label.
     * @return the generated key.
     * @since 1.0.0
     */
    protected static String generateKey(String label) {
        try {
            KeyGenerator keygen = KeyGenerator.getInstance(label, BouncyCastleProvider.PROVIDER_NAME);
            keygen.init(KEY_SIZE_MAP.get(label));
            SecretKey originalKey = keygen.generateKey();
            return bytesToHex(originalKey.getEncoded());
        } catch (Exception e) {
            LOGGER.error("generate key error", e);
        }
        return null;
    }

    /**
     * Encrypt or decrypt the given content by password and key label.
     *
     * @param content       the content to encrypt or decrypt.
     * @param password      the password for encrypt or decrypt.
     * @param forEncryption true is to encrypt false is to decrypt.
     * @param label         the key label.
     * @return encrypt or decrypt result.
     * @since 1.0.0
     */
    private static String doEncrypt(String content, String password, boolean forEncryption, String label) {
        if (content == null || password == null) {
            return null;
        }
        AEADBlockCipher cipher = initCipher(password, forEncryption, label);
        if (cipher == null) {
            return null;
        }
        byte[] data = forEncryption ? content.getBytes(StandardCharsets.UTF_8) : Base64.getMimeDecoder().decode(content);
        try (ByteArrayInputStream inputStream = new ByteArrayInputStream(data);
             CipherInputStream cipherInputStream = new CipherInputStream(inputStream, cipher);
             ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) {
            byte[] cache = new byte[BUFFER_SIZE];
            int r;
            while ((r = cipherInputStream.read(cache)) != -1) {
                outputStream.write(cache, 0, r);
                outputStream.flush();
            }
            byte[] result = outputStream.toByteArray();
            return forEncryption ? Base64.getEncoder().encodeToString(result) : new String(result, StandardCharsets.UTF_8);
        } catch (Exception e) {
            LOGGER.error("do encrypt error", e);
        }
        return null;
    }

    /**
     * Encrypt or decrypt the given file by password and key label.
     *
     * @param originFile    the origin file to encrypt or decrypt.
     * @param targetFile    the target path to store encrypt or decrypt file.
     * @param password      the password for encrypt or decrypt.
     * @param forEncryption true is to encrypt false is to decrypt.
     * @param label         the key label.
     * @return whether successful.
     * @since 1.0.0
     */
    private static boolean doEncryptFile(File originFile, File targetFile, String password, boolean forEncryption, String label) {
        if (originFile == null || !Files.isReadable(originFile.toPath()) || password == null || label == null) {
            LOGGER.error("do encrypt file param check failed");
            return false;
        }
        AEADBlockCipher cipher = initCipher(password, forEncryption, label);
        if (cipher == null) {
            return false;
        }
        try (InputStream inputStream = new FileInputStream(originFile);
             CipherInputStream cipherInputStream = new CipherInputStream(inputStream, cipher);
             OutputStream outputStream = new FileOutputStream(targetFile)) {
            byte[] cache = new byte[BUFFER_SIZE];
            int r;
            while ((r = cipherInputStream.read(cache)) != -1) {
                outputStream.write(cache, 0, r);
                outputStream.flush();
            }
            return true;
        } catch (Exception e) {
            LOGGER.error("do encrypt file error", e);
        }
        return false;
    }

    /**
     * Initialize the cipher by password and key label.
     *
     * @param password      the password.
     * @param forEncryption true is to encrypt false is to decrypt.
     * @param label         the key label.
     * @return {@link AEADBlockCipher} instance.
     * @since 1.0.0
     */
    private static AEADBlockCipher initCipher(String password, boolean forEncryption, String label) {
        try {
            AEADBlockCipher cipher = new GCMBlockCipher(BLOCK_CIPHER_MAP.get(label));
            cipher.init(forEncryption, new AEADParameters(new KeyParameter(hexToBytes(password)), MAC_SIZE, IV));
            return cipher;
        } catch (Exception e) {
            LOGGER.error("init cipher error", e);
        }
        return null;
    }

    /**
     * Byte array to hex string.
     *
     * @param bytes byte array.
     * @return hex string.
     * @since 1.0.0
     */
    private static String bytesToHex(byte[] bytes) {
        if (bytes == null) {
            return null;
        }
        byte[] hexChars = new byte[bytes.length << 1];
        int index = 0;
        for (int v : bytes) {
            hexChars[index++] = HEX_ARRAY[v >>> FOUR & F];
            hexChars[index++] = HEX_ARRAY[v & F];
        }
        return new String(hexChars, StandardCharsets.UTF_8);
    }

    /**
     * Hex string to byte array.
     *
     * @param hex hex string.
     * @return byte array.
     * @since 1.0.0
     */
    private static byte[] hexToBytes(String hex) {
        int byteLength = hex.length() >> 1;
        byte[] result = new byte[byteLength];
        for (int i = 0; i < byteLength; i++) {
            int offset = i << 1;
            result[i] = (byte) Integer.parseInt(hex.substring(offset, offset + TWO), SIXTEEN);
        }
        return result;
    }
}
